from iFinDPy import *
import pandas as pd
import numpy as np


# 登录函数
def thslogindemo():
    # 输入用户的帐号和密码
    thsLogin = THS_iFinDLogin("xnyf018", "SWSC600369")
    print(thsLogin)
    if thsLogin != 0:
        print('登录失败')
    else:
        print('登录成功')


thslogindemo()

data_g = THS_HQ('399370.SZ', 'close', '', '2013-12-31', '2024-07-31')

data_v = THS_HQ('399371.SZ', 'close', '', '2013-12-31', '2024-07-31')

close_g = data_g.data
close_v = data_v.data

# 确保数据按日期排序
close_g['time'] = pd.to_datetime(close_g['time'])
close_v['time'] = pd.to_datetime(close_v['time'])
close_g.sort_values(by='time', inplace=True)
close_v.sort_values(by='time', inplace=True)

close_g['dateeom'] = 0
close_g.loc[close_g.groupby([close_g['time'].dt.year, close_g['time'].dt.month])['time'].idxmax(), 'dateeom'] = 1

close_v['dateeom'] = 0
close_v.loc[close_v.groupby([close_v['time'].dt.year, close_v['time'].dt.month])['time'].idxmax(), 'dateeom'] = 1

close_g.set_index('time', inplace=True)
close_v.set_index('time', inplace=True)

print(type(close_g))
print(close_g.head())
print(type(close_v))
print(close_v.head())


# 年化收益率
def annualized_return(cumulative_returns, days_per_year=252):
    total_return = cumulative_returns.iloc[-1]
    num_days = len(cumulative_returns)
    return (total_return ** (days_per_year / num_days)) - 1


# 年化波动率
def annualized_volatility(daily_returns, days_per_year=252):
    return daily_returns.std() * np.sqrt(days_per_year)


# 收益风险比
def sharpe_ratio(annual_return, annual_volatility):
    return annual_return / annual_volatility


# 最大回撤
def max_drawdown(cumulative_returns):
    peak = cumulative_returns.cummax()
    drawdown = (cumulative_returns - peak) / peak
    return drawdown.min()


# 年均双边换手率
def turnover_rate(positions):
    turnover = positions.diff().abs().sum() * 2 / (len(positions) / 12)
    return turnover

# 计算21日收益率
close_g['return_21d'] = close_g['close'].pct_change(21)
close_v['return_21d'] = close_v['close'].pct_change(21)

# 每个月底的日期
month_ends = close_g[close_g['dateeom'] == 1].index

# 生成交易信号：如果成长相对收益高，则持有成长；否则持有价值
signals = pd.DataFrame(index=month_ends)
signals['g_vs_v'] = close_g['return_21d'].loc[month_ends] - close_v['return_21d'].loc[month_ends]
signals['position'] = signals['g_vs_v'].apply(lambda x: 1 if x > 0 else -1)  # 1代表成长，-1代表价值
signals.to_csv('signals.csv')

# 用于每个交易日的持仓
daily_signals = pd.DataFrame(index=close_g.index)
daily_signals['position'] = signals['position'].reindex(close_g.index, method='ffill')

# 计算策略的日收益
strategy_returns = pd.DataFrame(index=close_g.index)
strategy_returns['strategy'] = np.where(
    daily_signals['position'] == 1,
    close_g['close'].pct_change(),
    close_v['close'].pct_change()
)

strategy_returns.to_csv('strategy_returns.csv')

# 计算基准组合收益
benchmark_returns = (close_g['close'].pct_change() + close_v['close'].pct_change()) / 2

# 计算累计收益
strategy_cumulative = (1 + strategy_returns['strategy']).cumprod()
strategy_cumulative.to_csv('strategy_cumulative.csv')
benchmark_cumulative = (1 + benchmark_returns).cumprod()


# 计算策略和基准的业绩指标
strategy_annual_return = annualized_return(strategy_cumulative)
strategy_annual_volatility = annualized_volatility(strategy_returns['strategy'])
strategy_sharpe_ratio = sharpe_ratio(strategy_annual_return, strategy_annual_volatility)
strategy_max_drawdown = max_drawdown(strategy_cumulative)
strategy_turnover_rate = turnover_rate(signals['position'])

benchmark_annual_return = annualized_return(benchmark_cumulative)
benchmark_annual_volatility = annualized_volatility(benchmark_returns)
benchmark_sharpe_ratio = sharpe_ratio(benchmark_annual_return, benchmark_annual_volatility)
benchmark_max_drawdown = max_drawdown(benchmark_cumulative)
benchmark_turnover_rate = turnover_rate((close_g['close'].pct_change() + close_v['close'].pct_change()) / 2)

# 输出结果
results = {
    '指标': ['年化收益率', '年化波动率', '收益风险比', '最大回撤', '年均双边换手率'],
    '策略': [strategy_annual_return, strategy_annual_volatility, strategy_sharpe_ratio, strategy_max_drawdown, strategy_turnover_rate],
    '基准': [benchmark_annual_return, benchmark_annual_volatility, benchmark_sharpe_ratio, benchmark_max_drawdown, benchmark_turnover_rate]
}

results_df = pd.DataFrame(results)


print(results_df)

import matplotlib.pyplot as plt

# 绘制策略和基准的累计收益曲线
plt.figure(figsize=(14, 7))
plt.plot(strategy_cumulative.index, strategy_cumulative, label='策略累计收益', linestyle='-', linewidth=2)
plt.plot(benchmark_cumulative.index, benchmark_cumulative, label='基准累计收益', linestyle='--', linewidth=2)
plt.title('策略与基准的累计收益曲线')
plt.xlabel('日期')
plt.ylabel('累计收益')
plt.legend()
plt.grid(True)
plt.show()

